import logging
import os
import threading
import time
from typing import Callable, Dict, List, Tuple

import orjson as json

from tcp_comm.common import *
from tcp_comm.lock import RWLockable, TimeRWLock
from tcp_comm.message import *
from tcp_comm.tcp import TCPClient, TCPConnectionError

logger = logging.getLogger(__name__)


class Frame:
    def __init__(self, timestamps: float):
        self.timestamps = timestamps
        self.objects: Dict[str, ObjectInfo] = {}
        self.perceptions: Dict[str, List[Vector3D]] = {}

    def parse_message(self, msg: MessageBase):
        if isinstance(msg, ObjectInfoMsg):
            self.objects[msg.objectInfo.objectName] = msg.objectInfo
        elif isinstance(msg, UAVInfoMsg):
            self.objects[msg.objectInfo.objectName] = msg.objectInfo
            self.perceptions[msg.objectInfo.objectName] = msg.obstacles_pos
        else:
            logger.warning("Unhandled message type: %s", type(msg).__name__)

    def get_object(self, object_name: str):
        return self.objects.get(object_name, None)

    def get_object_perception(self, object_name: str):
        return self.perceptions.get(object_name, None)

    def to_dict(self):
        return {
            "timestamps": self.timestamps,
            "objects": self.objects,
            "perceptions": self.perceptions,
        }


class CommClient:
    def __init__(self, config: dict = None) -> None:
        self.config = config or {}
        host = self.config.get("host", "localhost")
        world_port = self.config.get("port", self.config.get("world_port", 8100))
        control_port = self.config.get("control_port", 8200)
        stream_port = self.config.get("stream_port", 8300)
        timeout = self.config.get("timeout", 10)

        self._world_client = TCPClient(host, world_port, 60)
        """TCPClient for simulation management"""

        self._stream_client = TCPClient(host, stream_port, timeout)
        """TCPClient for episode data streaming"""

        self._control_client = TCPClient(host, control_port, timeout)
        """TCPClient for command sending"""

        self._log_dir = self.config.get("log_dir", None)
        self._time_wait = self.config.get("time_wait", None)
        self._is_episode_requested = False
        self._stream_thread = None
        self._lock = TimeRWLock(timeout)
        self._event_handlers = {}
        self._frames = [Frame(-1)]

        if self._log_dir:
            os.makedirs(self._log_dir, exist_ok=True)

    def clear_frames(self):
        with self._lock.gen_wlock():
            if self._log_dir is not None and len(self._frames) > 1:
                timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
                with open(f"{self._log_dir}/sim_data_{timestamp}.json", "w") as f:
                    self._frames.pop(0)
                    f.write(
                        json.dumps(
                            self._frames,
                            default=lambda x: x.to_dict(),
                            option=json.OPT_SERIALIZE_NUMPY,
                        ).decode("utf-8")
                    )
            self._frames = [Frame(-1)]

    def connected(self):
        return self._world_client.connected()

    def connect(self):
        self._world_client.connect()

    def disconnect(self):
        self._world_client.disconnect()
        self._stream_client.disconnect()
        self._control_client.disconnect()
        if self._stream_thread:
            self._stream_thread.join()
            self._stream_thread = None

    def start_episode(self, messages: List[MessageBase]):
        """Start an new episode with the given messages.

        Args:
            messages (list): List of messages to be sent to the server, usually object creation messages.
        """
        if not self._is_episode_requested:
            self._request_new_episode()

        try:
            self._control_client.connect()

            # Scene creation messages
            for msg in messages:
                self.send_message(self._control_client, msg, self._time_wait)

            # Start listening to the stream
            self._stream_client.connect()
            self._stream_thread = threading.Thread(
                target=self.recv_thread,
                args=(
                    self._stream_client,
                    self._frames,
                    self._event_handlers,
                    self._lock,
                ),
            )
            self._stream_thread.start()

            # Start message
            start_msg = TaskStatusInfoMsg(timestamps=time.time(), taskStatus=1)
            self.send_message(self._world_client, start_msg, self._time_wait)
        finally:
            self._is_episode_requested = False

    def end_episode(self) -> SimulationResultMsg:
        """End the current episode."""
        self._control_client.disconnect()
        self._stream_client.disconnect()
        self._stream_thread.join()

        msg = TaskStatusInfoMsg(timestamps=time.time(), taskStatus=2)
        self.send_message(self._world_client, msg, self._time_wait)

        data = self._world_client.read()
        if not data:
            raise RuntimeError("Failed to receive episode end result")

        recv_msg = SimulationResultMsg.from_json(data.decode("utf-8"))
        return recv_msg

    def send_control(self, msg: MessageBase):
        """Send a control message to the server."""
        self.send_message(self._control_client, msg, self._time_wait)

    def get_frame(self) -> Frame:
        """Return the last **complete** frame."""
        with self._lock.gen_rlock():
            frame = self._frames[-2] if len(self._frames) > 2 else None
        return frame

    def register_handler(self, name: str, event_type: EventType, callback: Callable):
        """Register a callback function to be called when an event of the given type is received.

        Args:
            name (str): Identifier of the callback.
            event_type (EventType): Event type to be listened.
            callback (Callable): Callback function to be called. The function should accept a single argument "event".
        """
        self._event_handlers[name] = (event_type, callback)

    def _request_new_episode(self):
        self._stream_client.disconnect()
        self._control_client.disconnect()
        self.clear_frames()

        msg = TaskStatusInfoMsg(timestamps=time.time(), taskStatus=3)
        self.send_message(self._world_client, msg, self._time_wait)

        # Wait for Ready Status
        data = self._world_client.read()
        if not data:
            raise RuntimeError("Failed to receive simulation status")

        recv_msg: TaskStatusInfoMsg = MessageBase.from_json(data)
        if recv_msg.taskStatus != 0:
            raise RuntimeError("Invalid status, expected 0, got", recv_msg.taskStatus)

        self._is_episode_requested = True

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()
        return True

    @staticmethod
    def send_message(
        client: TCPClient,
        msg: MessageBase,
        wait_time: float = 0.0,
        ignore_return: bool = False,
    ):
        if wait_time:
            time.sleep(wait_time)

        logger.debug(f"Send message: {type(msg)}")
        data = json.dumps(msg.to_dict(), option=json.OPT_SERIALIZE_NUMPY)
        client.write(data)

        # receive return msg
        if not ignore_return and isinstance(
            msg, (ObjectInitActionMsg, CreateGroupInfoMsg, SetGroupTargetInfoMsg)
        ):
            try:
                data = client.read()
                res = MessageBase.from_json(data)
                if res.msgType == MsgType.MsgType_StatusVerify:
                    logger.debug(f"{type(msg)} successfully sent")
                else:
                    logger.warning(f"Wrong type of return msg: {type(res)}")
            except TCPConnectionError:
                logger.error(f"{msg} does not receive return msg")

    @staticmethod
    def recv_thread(
        client: TCPClient,
        frames: List[Frame],
        handlers: Dict[str, Tuple[EventType, Callable]],
        lock: RWLockable,
    ):
        while client.connected():
            try:
                data = client.read()
            except TCPConnectionError:
                logger.info(f"{client._logprefix}connection closed")
                break

            msg = MessageBase.from_json(data)
            if isinstance(msg, EventInfoMsg):
                for callback in handlers.values():
                    event_type, handler = callback
                    if msg.eventType == event_type:
                        handler(msg)
            else:
                with lock.gen_wlock():
                    last_frame = frames[-1]
                    if msg.timestamps != last_frame.timestamps:
                        frames.append(Frame(msg.timestamps))
                        last_frame = frames[-1]
                    last_frame.parse_message(msg)
